package cn.dataling.ai.config;

import cn.dataling.ai.domain.AiChatMemory;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import cn.dataling.ai.mapper.AiChatMemoryMapper;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 自定义对话以及存储
 * SpringAI 自带了基于 JdbcTemplate 的 JDBC 的会话存储类 在 spring-ai-starter-model-chat-memory-jdbc 依赖里面 但是只支持 postgresql 和 mariadb
 */
@Component
public class JdbcChatMemory implements ChatMemory {

    @Resource
    private AiChatMemoryMapper aiChatMemoryMapper;

    @Override
    public void add(String conversationId, List<Message> messages) {

        List<AiChatMemory> aiChatMemories = messages.stream().map(e -> {
            AiChatMemory aiChatMemory = new AiChatMemory();
            aiChatMemory.setContent(e.getText());
            aiChatMemory.setConversationId(conversationId);
            aiChatMemory.setCreatedAt(LocalDateTime.now());
            aiChatMemory.setType(e.getMessageType().name());
            return aiChatMemory;
        }).collect(Collectors.toList());

        aiChatMemoryMapper.insert(aiChatMemories);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {

        List<AiChatMemory> aiChatMemories = aiChatMemoryMapper.selectList(new LambdaQueryWrapper<AiChatMemory>()
                .eq(AiChatMemory::getConversationId, conversationId)
                .orderByDesc(AiChatMemory::getCreatedAt)
                .last("LIMIT " + lastN)
        );

        return aiChatMemories.stream().map(e -> {
            MessageType messageType = MessageType.valueOf(e.getType());
            return switch (messageType) {
                case USER -> new UserMessage(e.getContent());
                case ASSISTANT -> new AssistantMessage(e.getContent());
                case SYSTEM -> new SystemMessage(e.getContent());
                default -> null;
            };
        }).collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        aiChatMemoryMapper.delete(new LambdaQueryWrapper<AiChatMemory>()
                .eq(AiChatMemory::getConversationId, conversationId));
    }
}
